import rasterio
import numpy as np
import pandas as pd
from tqdm import tqdm

def read_map(file_path):
    with rasterio.open(file_path) as src:
        elevation_map = src.read(1).astype(np.float64)
        return elevation_map

def calculate_slope_and_aspect(elevation_map, cell_size=5):
    rows, cols = elevation_map.shape
    slope_map = np.full((rows, cols), -1.0, dtype=np.float64)
    aspect_map = np.full((rows, cols), -1.0, dtype=np.float64)
    a = elevation_map[:-2, 2:]
    b = elevation_map[1:-1, 2:]
    c = elevation_map[2:, 2:]
    d = elevation_map[:-2, 1:-1]
    f = elevation_map[2:, 1:-1]
    g = elevation_map[:-2, :-2]
    h = elevation_map[1:-1, :-2]
    i = elevation_map[2:, :-2]
    dz_dx = (c + 2*f + i - a - 2*d - g) / (8 * cell_size)*5
    dz_dy = (a + 2*b + c - g - 2*h - i) / (8 * cell_size)*5
    slope = np.arctan(np.sqrt(dz_dx**2 + dz_dy**2))
    slope_deg = np.degrees(slope)
    
    # 计算坡向（地理坐标：0°为正北，顺时针增加）
    aspect = np.zeros_like(slope_deg)
    result = (270 + np.arctan2(dz_dx, dz_dy) * (180 / np.pi) - 90 * np.sign(dz_dy)) % 360
    aspect[(dz_dx != 0) & (dz_dy != 0)] = result[(dz_dx != 0) & (dz_dy != 0)]
    aspect[(dz_dy == 0) & (dz_dx < 0)] = 90
    aspect[(dz_dy == 0) & (dz_dx > 0)] = 270
    aspect[(dz_dx == 0) & (dz_dy > 0)] = 180
    aspect[(dz_dx == 0) & (dz_dy < 0)] = 0
    aspect[(dz_dx == 0) & (dz_dy == 0)] = -1

    slope_map[1:-1, 1:-1] = slope_deg
    aspect_map[1:-1, 1:-1] = aspect
    return slope_map, aspect_map

def build(map_path="map.tif"):
    elevation_map = read_map(map_path)
    rows, cols = elevation_map.shape
    M = np.zeros((rows, cols, 4))
    array_index_y = 12499 - np.arange(rows)
    array_index_x = np.arange(cols)
    Y, X = np.meshgrid(array_index_y, array_index_x, indexing='ij')
    M[X, Y, 0] = elevation_map
    slope, aspect = calculate_slope_and_aspect(M[:, :, 0])
    M[:, :, 1] = slope
    M[:, :, 2] = aspect
    zone_path = '附件4：不良区域位置信息.xlsx'
    xls = pd.ExcelFile(zone_path)
    
    # 读取所有子表
    all_coordinates = []
    for sheet_name in xls.sheet_names:
        df = pd.read_excel(xls, sheet_name=sheet_name)
        coordinates = df[['栅格x坐标', '栅格y坐标']]
        all_coordinates.extend(coordinates.values.tolist())

    for i in all_coordinates:
        M[i[0], i[1], 3] = 1

    np.save('map_data.npy', M)

if __name__=="__main__":
    build()